import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import numpy as np

# 配置设备（优先用 GPU）
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# ------------------- 1. 数据加载与预处理 ------------------- #
# 通用转换：转 Tensor + 归一化
transform_fashion = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))  # FashionMNIST 是单通道
])

transform_cifar = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))  # CIFAR-10 是 3 通道
])

# 加载 FashionMNIST 数据集
train_dataset_fashion = torchvision.datasets.FashionMNIST(
    root='./data', train=True, download=True, transform=transform_fashion
)
test_dataset_fashion = torchvision.datasets.FashionMNIST(
    root='./data', train=False, download=True, transform=transform_fashion
)

# 加载 CIFAR-10 数据集
train_dataset_cifar = torchvision.datasets.CIFAR10(
    root='./data', train=True, download=True, transform=transform_cifar
)
test_dataset_cifar = torchvision.datasets.CIFAR10(
    root='./data', train=False, download=True, transform=transform_cifar
)

# 数据加载器（batch 训练）
batch_size = 64
train_loader_fashion = torch.utils.data.DataLoader(
    train_dataset_fashion, batch_size=batch_size, shuffle=True
)
test_loader_fashion = torch.utils.data.DataLoader(
    test_dataset_fashion, batch_size=batch_size, shuffle=False
)

train_loader_cifar = torch.utils.data.DataLoader(
    train_dataset_cifar, batch_size=batch_size, shuffle=True
)
test_loader_cifar = torch.utils.data.DataLoader(
    test_dataset_cifar, batch_size=batch_size, shuffle=False
)

# ------------------- 2. 构建 3 层神经网络 ------------------- #
class ThreeLayerNet(nn.Module):
    def __init__(self, input_size, hidden_size, num_classes, is_cifar=False):
        super(ThreeLayerNet, self).__init__()
        self.is_cifar = is_cifar  # 区分单通道（FashionMNIST）和 3 通道（CIFAR-10）
        if is_cifar:
            input_size = 3 * 32 * 32  # CIFAR-10 是 32×32×3
        else:
            input_size = 1 * 28 * 28  # FashionMNIST 是 28×28×1

        self.fc1 = nn.Linear(input_size, hidden_size)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_size, hidden_size // 2)  # 第二层隐藏层
        self.fc3 = nn.Linear(hidden_size // 2, num_classes)

    def forward(self, x):
        # 展平输入
        if self.is_cifar:
            x = x.view(-1, 3 * 32 * 32)
        else:
            x = x.view(-1, 1 * 28 * 28)

        out = self.fc1(x)
        out = self.relu(out)
        out = self.fc2(out)
        out = self.relu(out)
        out = self.fc3(out)
        return out

# 初始化模型（分别适配两个数据集）
# FashionMNIST 配置
fashion_input_size = 1 * 28 * 28
fashion_hidden_size = 256  # 可调整
fashion_num_classes = 10
model_fashion = ThreeLayerNet(fashion_input_size, fashion_hidden_size, fashion_num_classes, is_cifar=False).to(device)

# CIFAR-10 配置
cifar_input_size = 3 * 32 * 32
cifar_hidden_size = 512  # 可调整
cifar_num_classes = 10
model_cifar = ThreeLayerNet(cifar_input_size, cifar_hidden_size, cifar_num_classes, is_cifar=True).to(device)

# ------------------- 3. 训练与评估函数（通用） ------------------- #
def train_evaluate(model, train_loader, test_loader, num_epochs, dataset_name):
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

    train_losses = []
    train_accuracies = []
    test_losses = []
    test_accuracies = []

    for epoch in range(num_epochs):
        # 训练阶段
        model.train()
        train_correct = 0
        train_total = 0
        train_running_loss = 0.0

        for images, labels in train_loader:
            images, labels = images.to(device), labels.to(device)
            outputs = model(images)
            loss = criterion(outputs, labels)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            train_running_loss += loss.item()
            _, predicted = torch.max(outputs.data, 1)
            train_total += labels.size(0)
            train_correct += (predicted == labels).sum().item()

        train_epoch_loss = train_running_loss / len(train_loader)
        train_epoch_accuracy = 100 * train_correct / train_total
        train_losses.append(train_epoch_loss)
        train_accuracies.append(train_epoch_accuracy)

        # 测试阶段
        model.eval()
        test_correct = 0
        test_total = 0
        test_running_loss = 0.0

        with torch.no_grad():
            for images, labels in test_loader:
                images, labels = images.to(device), labels.to(device)
                outputs = model(images)
                loss = criterion(outputs, labels)

                test_running_loss += loss.item()
                _, predicted = torch.max(outputs.data, 1)
                test_total += labels.size(0)
                test_correct += (predicted == labels).sum().item()

        test_epoch_loss = test_running_loss / len(test_loader)
        test_epoch_accuracy = 100 * test_correct / test_total
        test_losses.append(test_epoch_loss)
        test_accuracies.append(test_epoch_accuracy)

        # 打印日志
        print(f'{dataset_name} | Epoch {epoch + 1}/{num_epochs}')
        print(f'Train Loss: {train_epoch_loss:.4f} | Train Acc: {train_epoch_accuracy:.2f}%')
        print(f'Test Loss: {test_epoch_loss:.4f} | Test Acc: {test_epoch_accuracy:.2f}%')
        print('-' * 30)

    # 可视化曲线
    visualize_curves(train_losses, test_losses, train_accuracies, test_accuracies, dataset_name)
    # 可视化预测结果
    visualize_predictions(model, test_loader, dataset_name, is_cifar=(dataset_name == 'CIFAR-10'))

    return test_epoch_accuracy

# ------------------- 4. 可视化工具函数 ------------------- #
def visualize_curves(train_losses, test_losses, train_accuracies, test_accuracies, dataset_name):
    plt.figure(figsize=(12, 5))

    # Loss 曲线
    plt.subplot(1, 2, 1)
    plt.plot(train_losses, label='Train Loss')
    plt.plot(test_losses, label='Test Loss')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.title(f'{dataset_name} - Training vs Test Loss')
    plt.legend()

    # Accuracy 曲线
    plt.subplot(1, 2, 2)
    plt.plot(train_accuracies, label='Train Accuracy')
    plt.plot(test_accuracies, label='Test Accuracy')
    plt.xlabel('Epoch')
    plt.ylabel('Accuracy (%)')
    plt.title(f'{dataset_name} - Training vs Test Accuracy')
    plt.legend()

    plt.show()

def visualize_predictions(model, test_loader, dataset_name, is_cifar=False):
    model.eval()
    with torch.no_grad():
        # 取一个 batch 的测试数据
        images, labels = next(iter(test_loader))
        images, labels = images.to(device), labels.to(device)
        outputs = model(images)
        _, predicted = torch.max(outputs, 1)

        # 转换为 numpy（方便可视化）
        images = images.cpu().numpy()
        labels = labels.cpu().numpy()
        predicted = predicted.cpu().numpy()

        # 还原归一化（方便显示图片）
        if is_cifar:
            images = images * 0.5 + 0.5  # CIFAR-10 是 3 通道
            num_show = 10
            fig, axes = plt.subplots(2, 5, figsize=(15, 6))
        else:
            images = images[:, 0, :, :] * 0.5 + 0.5  # FashionMNIST 是单通道
            num_show = 10
            fig, axes = plt.subplots(2, 5, figsize=(15, 6))

        axes = axes.ravel()
        for i in range(num_show):
            if is_cifar:
                img = np.transpose(images[i], (1, 2, 0))  # CIFAR-10 需调整通道维度
            else:
                img = images[i]

            axes[i].imshow(img, cmap='gray' if not is_cifar else None)
            axes[i].set_title(f'True: {labels[i]}\nPred: {predicted[i]}')
            axes[i].axis('off')

        plt.suptitle(f'{dataset_name} - Prediction Examples')
        plt.show()

# ------------------- 5. 主流程：训练两个数据集 ------------------- #
if __name__ == "__main__":
    num_epochs = 10  # 可调整，建议先小试再增大

    # 训练 FashionMNIST
    print('=== FashionMNIST 训练 ===')
    fashion_test_acc = train_evaluate(
        model_fashion, train_loader_fashion, test_loader_fashion,
        num_epochs, 'FashionMNIST'
    )
    print(f'FashionMNIST 测试集最终准确率: {fashion_test_acc:.2f}%\n')

    # 训练 CIFAR-10
    print('=== CIFAR-10 训练 ===')
    cifar_test_acc = train_evaluate(
        model_cifar, train_loader_cifar, test_loader_cifar,
        num_epochs, 'CIFAR-10'
    )
    print(f'CIFAR-10 测试集最终准确率: {cifar_test_acc:.2f}%')